import numpy as np
import casadi as ca

def dh_transform(theta, d, a, alpha):
    ct = ca.cos(theta)
    st = ca.sin(theta)
    ca_ = ca.cos(alpha)
    sa_ = ca.sin(alpha)
    return ca.vertcat(
        ca.horzcat(ct, -st*ca_, st*sa_, a*ct),
        ca.horzcat(st, ct*ca_, -ct*sa_, a*st),
        ca.horzcat(0, sa_, ca_, d),
        ca.horzcat(0, 0, 0, 1)
    )

def ur5_fk_position_rotation(q):
    """
    :param q: CasADi SX/SX (6,)
    :return: position (3,), rotation matrix (3x3)
    """
    d = [0.0892, 0, 0, 0.10915, 0.09465, 0.0823]
    a = [0, -0.425, -0.3922, 0, 0, 0]
    alpha = [np.pi/2, 0, 0, np.pi/2, -np.pi/2, 0]

    T = ca.SX.eye(4)
    for i in range(6):
        T = T @ dh_transform(q[i], d[i], a[i], alpha[i])

    position = T[0:3, 3]
    rotation = T[0:3, 0:3]

    return position, rotation

def make_fk_function():
    q = ca.SX.sym("q", 6)
    pos, R = ur5_fk_position_rotation(q)
    fk_func = ca.Function("fk", [q], [pos, R])
    return fk_func